# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2024/3/8
"""
对比不同库中hann窗函数的实现
如果对称(sym=True)的话，有两个最大值，如果不对称(sym=False)的话，有一个最大值

- numpy的hanning函数是对称的
- scipy有hanning函数有sym参数设置，默认是对称的
- torch的hanning函数有periodic参数设置，默认是非对称的
"""
import numpy as np

import torch
import scipy.signal as signal

window_len = 512


def hann_sym(window_len):
    """对称hann窗"""
    win = np.zeros(window_len)
    for i in range(window_len):
        win[i] = 0.5 - 0.5 * np.cos(2 * np.pi * i / (window_len - 1))
    return win


def hann_asym(window_len):
    """非对称hann"""
    p_win = np.zeros(window_len)
    for i in range(window_len):
        p_win[i] = np.sin(np.pi * i / window_len)
        p_win[i] = p_win[i] * p_win[i]
    return p_win


def my_hann_aysm(win_len):
    haning_window = np.hanning(win_len + 1)  # 对称的hann窗
    out = haning_window[0:-1].astype('float32')  # 舍弃最后一个元素
    return out


scipy_sym = signal.windows.hann(window_len, sym=True)  # 对称的hann窗
scipy_Asym = signal.windows.hann(window_len, sym=False)  # 非对称的hann窗
hann_sym_c = hann_sym(window_len)
hann_asym_c = hann_asym(window_len)
my_hann= my_hann_aysm(window_len)

print(np.allclose(scipy_sym, hann_sym_c))  # True
print(np.allclose(scipy_Asym, hann_asym_c))  # True
print(np.allclose(my_hann, hann_asym_c))  # True

numpy_window = np.hanning(window_len)  # 说明numpy的hanning函数是对称的
print(np.allclose(numpy_window, scipy_sym))  # True

torch_window = torch.hann_window(window_len)  # 非对称
torch_window_periodic = torch.hann_window(window_len, periodic=False)  # 非周期=对称
# print(torch.argmax(window_torch))

# 判断两个窗函数是否相等
print(np.allclose(scipy_Asym, torch_window.numpy(), rtol=1e-3))  # True
print(np.allclose(scipy_sym, torch_window_periodic.numpy(), rtol=1e-3))  # True
